package com.senke.aiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 查询重写器
 */
@Component
public class MyQueryRewriter {

    private final RewriteQueryTransformer rewriteQueryTransformer;

    public MyQueryRewriter(ChatModel dashscopeChatModel) {
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        RewriteQueryTransformer rewriteQueryTransformer = RewriteQueryTransformer
                .builder()
                .chatClientBuilder(builder)
                .build();
        this.rewriteQueryTransformer = rewriteQueryTransformer;
    }

    public String doQueryRewriter(String prompt) {
        Query transformedQuery = rewriteQueryTransformer.transform(new Query(prompt));
        return transformedQuery.text();
    }
}
